#include <stdlib.h>

#include "env.h"
#include "errmsg.h"
#include "semantic.h"
#include "types.h"

static table_t _venv;
static table_t _tenv;
static type_t this_class;

static void trans_decl(ast_decl_t decl);
static void trans_stmt(ast_stmt_t stmt);
static type_t trans_expr(ast_expr_t expr);
static type_t trans_type(ast_type_t type);

static void trans_method_decl(ast_decl_t decl);
static void trans_var_decl(ast_decl_t decl);
static type_t trans_class_type(ast_type_t type);

static type_t lookup_type(symbol_t name, int pos)
{
	type_t type = sym_lookup(_tenv, name);
	if (type)
		type = ty_actual(type);
	return type;
}

/*
  ####################################################################
  trans decl
  ####################################################################
*/

static void trans_class_decl(ast_decl_t decl)
{
	list_t p, q;
	
	// check for re-decls of class
	this_class = trans_type(decl->u.class_.name);
	if (this_class != NULL)   // class already defined
	  em_error(decl->pos, "class redefined");

	// create a new class type
	this_class = ty_class();
	
	// add new class type to symbol table
	sym_enter(_tenv, (decl->u.class_.name)->u.class_, this_class);

	// check for re-decls of var
	for (p = decl->u.class_.var_decls;
			p && p->next; p = p->next)
		for (q = p->next; q; q = q->next)
		{
			ast_decl_t vd1 = (ast_decl_t) p->data;
			ast_decl_t vd2 = (ast_decl_t) q->data;
			if (sym_name(vd1->u.var.name) == 
				sym_name(vd2->u.var.name))
				em_error((vd2->u.var.type)->pos, "'%s': variable redefined", sym_name(vd2->u.var.name));
		}

	// enter class scope
	sym_begin_scope(_venv);
	
	// trans var decls
	for (p = decl->u.class_.var_decls; p; p = p->next)
	{
		trans_var_decl((ast_decl_t)p->data);
	}

	// check for re-decls of methods
	for (p = decl->u.class_.method_decls;
			p && p->next; p = p->next)
		for (q = p->next; q; q = q->next)
		{
			ast_decl_t md1 = (ast_decl_t) p->data;
			ast_decl_t md2 = (ast_decl_t) q->data;
			if (sym_name(md1->u.method.name) == 
				sym_name(md2->u.method.name))
				em_error(md2->pos, "'%s': method redefined", sym_name(md2->u.method.name));
		}

	for (p = decl->u.class_.method_decls; p; p = p->next)
	{
		trans_method_decl((ast_decl_t) p->data);
	}
  
  this_class = NULL;
  
	// exit class scope
	sym_end_scope(_venv);
}

static void trans_var_decl(ast_decl_t decl)
{
  // get var type
  type_t type = trans_type(decl->u.var.type);
  if (type == NULL) {
    em_error(decl->pos, "'%s': undefined type", sym_name((decl->u.var.type)->u.class_));
    return;
  }
  // add var to env
	sym_enter(_venv, decl->u.var.name, env_var_entry(type));
}

static void trans_method_decl(ast_decl_t decl)
{
  // get formals list
  list_t p, q, formals = NULL;
  for (p = decl->u.method.para_decls; p; p = p->next)
  {
    ast_decl_t para = p->data;
    type_t type = trans_type(para->u.para.type);
    if (type == NULL) {
      em_error((para->u.para.type)->pos, "'%s': undefined type", sym_name((para->u.para.type)->u.class_));
      return;
    }
    else {
      formals = list(type, formals);
    }
  }
  
  // debug
  /*
  for (p = formals; p; p = p->next)
  {
    switch (((type_t)p->data)->kind) {
    case TY_CLASS:
      printf("class ");
      break;
    case TY_INT:
      printf("int ");
      break;
    case TY_FLOAT:
      printf("float ");
      break;
    case TY_ARRAY:
      printf("array %d ", ((type_t)p->data)->u.array->kind);
      break;
    }
    printf("\n");
  }
  */
  
  // get result type
  type_t result;
  if (decl->u.method.type == NULL) // void
	  result = ty_void();
  else result = trans_type(decl->u.method.type);
  if (result == NULL) {
	  em_error(decl->pos, "'%s': return type undefined", sym_name((decl->u.method.type)->u.class_));
	  return;
  }
  
  // although we can sure that right now this method is semantically right
  // but to support recursive method call, we need to add it to env first
  // create an env_func_entry & add func to _venv
  sym_enter(_venv, decl->u.method.name, env_func_entry(formals, result));
  
  // add method descriptors to 'this_class'
  sym_enter(this_class->u.methods, decl->u.method.name, env_func_entry(formals, result));
  
  // enter method scope
  sym_begin_scope(_venv);
  // check for duplicate param names
  for (p = decl->u.method.para_decls; p && p->next; p = p->next)
    for (q = p->next; q; q = q->next)
    {
      ast_decl_t para1 = (ast_decl_t) p->data;
      ast_decl_t para2 = (ast_decl_t) q->data;
      symbol_t para_name1 = para1->u.para.name;
      symbol_t para_name2 = para2->u.para.name;
      if ( sym_name(para_name1) == sym_name(para_name2) ) {
        em_error(para2->u.para.type->pos, "'%s': redefined parameter", sym_name(para_name2));
        return;
      }
    }
  // add params to _venv
  for (p = decl->u.method.para_decls; p; p = p->next)
  {
    ast_decl_t para = (ast_decl_t) p->data;
    type_t type = trans_type(para->u.para.type);
    if (type == NULL) {
      em_error(decl->pos, "'%s': undefined type", sym_name((decl->u.var.type)->u.class_));
      return;
    }
    //printf("add: %s\n", sym_name(para->u.para.name));
	  sym_enter(_venv, para->u.para.name, env_var_entry(type));
	}
  // trans_var_decl (s)
  // check for re-decls of var
  for (p = decl->u.method.var_decls; p && p->next; p = p->next)
  {
    ast_decl_t vd1 = (ast_decl_t) p->data;
    for (q = p->next; q; q = q->next)
    {
			ast_decl_t vd2 = (ast_decl_t) q->data;
			if (sym_name(vd1->u.var.name) == 
				sym_name(vd2->u.var.name))
				em_error((vd2->u.var.type)->pos, "'%s': variable redefined", sym_name(vd2->u.var.name));
		}
	}
	
  // add vars to _venv
  for (p = decl->u.method.var_decls; p; p = p->next)
  {
    ast_decl_t vd = (ast_decl_t) p->data;
    
    void *t = sym_lookup(_venv, vd->u.var.name);
	  //printf("find: %s\n", sym_name(vd->u.var.name));
	  if (t != NULL) 
	    em_error((vd->u.var.type)->pos, "Warning: variable '%s' will shadow outer variables", sym_name(vd->u.var.name));

    type_t type = trans_type(vd->u.var.type);
    if (type == NULL) {
      em_error(vd->pos, "'%s': undefined type", sym_name((vd->u.var.type)->u.class_));
      return;
    }
	  sym_enter(_venv, vd->u.method.name, env_var_entry(type));
	}

  // trans_stmt_decl (s)
  for (p = decl->u.method.stmts; p; p = p->next)
    trans_stmt((ast_stmt_t)p->data);

  // handle return stmt
  // check func return type matches the result type
  if (decl->u.method.ret == NULL) {
    if (result->kind != TY_VOID) {
      em_error(decl->pos, "return type don't match the result type");
      return;
    }
  }
  else { 
    type_t ret = trans_expr(decl->u.method.ret);
    if (ret == NULL || ret->kind != result->kind) {
      em_error(decl->pos, "return type don't match the result type");
      return;
    }
  }
  
  // exit method scope
  sym_end_scope(_venv);
}

static void trans_method_decl_stub(ast_decl_t decl)
{}

static void trans_para_decl_stub(ast_decl_t decl)
{}

typedef void (*trans_decl_func)(ast_decl_t);

static trans_decl_func _trans_decl_funcs[]
=
{
	trans_class_decl,
	trans_var_decl,
	trans_method_decl_stub,
	trans_para_decl_stub,
};

static void trans_decl(ast_decl_t decl)
{
	_trans_decl_funcs[decl->kind](decl);
}

/*
  ####################################################################
  trans stmt
  ####################################################################
*/
void trans_stmts_stmt(ast_stmt_t stmt)
{
  list_t p;
  for (p = stmt->u.stmts; p; p = p->next)
    trans_stmt((ast_stmt_t) p->data);
}

void trans_if_stmt(ast_stmt_t stmt)
{
  type_t type = trans_expr(stmt->u.if_.cond);
  if (type != NULL && type->kind != TY_BOOLEAN)
    em_error(stmt->pos, "if condition's type must be boolean");
  // trans if block
  // assert( stmt->u.if_.if_block != NULL );
  trans_stmt(stmt->u.if_.if_block);
  // trans else block
  if (stmt->u.if_.else_block != NULL)
    trans_stmt(stmt->u.if_.else_block);
}

void trans_while_stmt(ast_stmt_t stmt)
{
  type_t type = trans_expr(stmt->u.while_.cond);
  if (type != NULL && type->kind != TY_BOOLEAN)
    em_error(stmt->pos, "while condition's type must be boolean");
  // trans block
  trans_stmt(stmt->u.while_.block);
}

void trans_var_assign_stmt(ast_stmt_t stmt)
{
  env_entry_t lhs = sym_lookup(_venv, stmt->u.var_assign.lhs);
  type_t rhs = trans_expr(stmt->u.var_assign.rhs);
  if (lhs == NULL)
    em_error(stmt->pos, "variable type in the left side of assignment is undefined");
  else if (rhs == NULL)
    em_error(stmt->pos, "expression type in the right side of assignment is undefined");
  else if (lhs->u.var.type->kind != rhs->kind) {
    //printf("lhs: %d\n", lhs->u.var.type->kind);
    //printf("rhs: %d\n", rhs->kind);
    em_error(stmt->pos, "type in both sides of assignment must be the same");
  }
}

void trans_arr_assign_stmt(ast_stmt_t stmt)
{
  env_entry_t lhs_id = sym_lookup(_venv, stmt->u.arr_assign.lhs_id);
  if (lhs_id == NULL || lhs_id->u.var.type->kind != TY_ARRAY) {
    em_error(stmt->pos, "variable type in the left side of array assignment must be array type");
    return;
  }
  type_t lhs_index = trans_expr(stmt->u.arr_assign.lhs_index);
  if (lhs_index == NULL || lhs_index->kind != TY_INT) {
    em_error(stmt->pos, "array index must be int type");
    return;
  }
  type_t rhs = trans_expr(stmt->u.arr_assign.rhs);
  if (rhs == NULL) {
    em_error(stmt->pos, "expression type in the right side of array assignment is undefined");
    return;
  }
  if (lhs_id->u.var.type->u.array->kind != rhs->kind)
    em_error(stmt->pos, "type in both sides of array assignment must be the same");
}

void trans_call_stmt(ast_stmt_t stmt)
{
  type_t field_type = trans_expr(stmt->u.call.field);
  if (field_type == NULL)
  {
    em_error(stmt->pos, "calling undefined type field");
    return;
  }
  
  if (field_type->kind != TY_CLASS)
  {
    em_error(stmt->pos, "calling non-class type field");
    return;
  }
  
  env_entry_t func = sym_lookup(field_type->u.methods, stmt->u.call.name);
  
  if (func == NULL) {
    em_error(stmt->pos, "calling undefined method");
    return;
  }
  
  list_t p, q, r = NULL;
  // trick to reverse the list int stmt call args
  for (q = stmt->u.call.args; q; q = q->next)
  {
    r = list(q->data, r);
  }
  
  for (p = func->u.func.formals;
        p && r; p = p->next, r = r->next)
  {
    type_t formal_type = p->data;
    type_t args_type = trans_expr(r->data);
    if (formal_type == NULL || args_type == NULL)
    {
      em_error(stmt->pos, "undefined parameter types referred");
      return;
    }
    else if (formal_type->kind != args_type->kind)
    {
      em_error(stmt->pos, "parameter types don't match function prototype");
      //printf("formals: %d, calls: %d\n", formal_type->kind, args_type->kind);
      return;
    }
  }
  if (p != NULL || r != NULL)
    em_error(stmt->pos, "parameter types don't match function prototype");
}

typedef void (*trans_stmt_func)(ast_stmt_t);

static trans_stmt_func _trans_stmt_funcs[]
=
{
  trans_stmts_stmt,
  trans_if_stmt,
  trans_while_stmt,
  trans_var_assign_stmt,
  trans_arr_assign_stmt,
  trans_call_stmt,
};

static void trans_stmt(ast_stmt_t stmt)
{
  _trans_stmt_funcs[stmt->kind](stmt);
}

/*
  ####################################################################
  trans expr
  ####################################################################
*/

type_t trans_op_expr(ast_expr_t expr)
{
  type_t lhs = trans_expr(expr->u.op.left);
  type_t rhs = trans_expr(expr->u.op.right);
  //printf("%d, %d\n", lhs->kind, rhs->kind);
  if (lhs == NULL || rhs == NULL)
    return NULL;
  if (lhs->kind != rhs->kind ||
        lhs->kind == TY_CLASS || lhs->kind == TY_ARRAY) {
    em_error(expr->pos, "operands must be the same primitive type");
    return NULL;
  }
  if (AST_AND <= expr->u.op.op && 
        expr->u.op.op <= AST_NEQ)
    return ty_boolean();
  else
    return lhs;
}

type_t trans_not_expr(ast_expr_t expr)
{
  type_t type = trans_expr(expr->u.not_);
  if (type->kind != TY_BOOLEAN) {
    em_error(expr->pos, "not operation can only applied on boolean type");
    return NULL;
  }
  return type;
}

type_t trans_arr_access_expr(ast_expr_t expr)
{
  type_t arr_type = trans_expr(expr->u.arr_access.arr);
  if (arr_type == NULL) {
    em_error(expr->pos, "undefined type");
    return NULL;
  }
  if (arr_type->kind != TY_ARRAY) {
    em_error(expr->pos, "not a array type");
    return NULL;
  }
  else return arr_type->u.array;
}

type_t trans_method_access_expr(ast_expr_t expr)
{
  type_t field_type = trans_expr(expr->u.method_access.field);
  if (field_type == NULL)
  {
    em_error(expr->pos, "invoking undefined type field");
    return NULL;
  }
  
  if (field_type->kind != TY_CLASS)
  {
    em_error(expr->pos, "invoking non-class type field");
    return NULL;
  }
  
  env_entry_t func = sym_lookup(field_type->u.methods, expr->u.method_access.name);
  list_t p, q, r = NULL;
  // trick to reverse the list int stmt call args
  for (q = expr->u.method_access.args; q; q = q->next)
  {
    r = list(q->data, r);
  }
  
  for (p = func->u.func.formals;
        p && r; p = p->next, r = r->next)
  {
    type_t formal_type = p->data;
    type_t args_type = trans_expr(r->data);
    if (formal_type == NULL || args_type == NULL)
    {
      em_error(expr->pos, "undefined parameter types referred");
      return NULL;
    }
    else if (formal_type->kind != args_type->kind)
    {
      em_error(expr->pos, "parameter types don't match function prototype");
      //printf("formals: %d, calls: %d\n", formal_type->kind, args_type->kind);
      return NULL;
    }
  }
  if (p != NULL || r != NULL) {
    em_error(expr->pos, "parameter types don't match function prototype");
    return NULL;
  }
  
  return func->u.func.result;
}

type_t trans_arr_length_expr(ast_expr_t expr)
{
  return NULL;
}

type_t trans_idf_expr(ast_expr_t expr)
{  
  env_entry_t t = sym_lookup(_venv, expr->u.idf);
  if (t == NULL) {
    em_error(expr->pos, "'%s': undefined variable", sym_name(expr->u.idf));
    return NULL;
  } else return t->u.var.type;
  //else printf("lookup idf: %s, type: %d\n", sym_name(expr->u.idf), t->u.var.type->kind);
}

type_t trans_this_expr(ast_expr_t expr)
{
  return this_class;
}

type_t trans_class_init_expr(ast_expr_t expr)
{
  return lookup_type(expr->u.class_init, expr->pos);
}

type_t trans_arr_init_expr(ast_expr_t expr)
{
  return trans_type(expr->u.arr_init.type);
  //TODO 
  // what about array size??
}

type_t trans_arr_decl_expr(ast_expr_t expr)
{
  list_t p = expr->u.arr_decl;
  ast_expr_t p_data = (ast_expr_t) p->data;
  
  type_t type = trans_expr(p_data);
  if (type == NULL) return NULL;
  
  // all values in array must be the same type
  for (p = expr->u.arr_decl; p; p = p->next)
  {
    type_t aType = trans_expr(p->data);
    if (type->kind != aType->kind) {
      em_error(expr->pos, "all values in array must be the same type");
      return NULL;
    }
  }
  return ty_array(type);
}

type_t trans_int_expr(ast_expr_t expr)
{
  return ty_int();
}

type_t trans_float_expr(ast_expr_t expr)
{
  return ty_float();
}

type_t trans_char_expr(ast_expr_t expr)
{
  return ty_char();
}

type_t trans_string_expr(ast_expr_t expr)
{
  return ty_string();
}

type_t trans_bool_expr(ast_expr_t expr)
{
  return ty_boolean();
}

typedef type_t (*trans_expr_func)(ast_expr_t);

static trans_expr_func _trans_expr_funcs[]
=
{
  trans_op_expr,
  trans_not_expr,
  trans_arr_access_expr,
  trans_method_access_expr,
  trans_arr_length_expr,
  trans_idf_expr,
  trans_this_expr,
  trans_class_init_expr,
  trans_arr_init_expr,
  trans_arr_decl_expr,
  trans_int_expr,
  trans_float_expr,
  trans_char_expr,
  trans_string_expr,
  trans_bool_expr,
};

static type_t trans_expr(ast_expr_t expr)
{
  //printf("%d\n", expr->kind);
  return _trans_expr_funcs[expr->kind](expr);
}

/*
  ####################################################################
  trans type
  ####################################################################
*/

static type_t trans_primitive_type(ast_type_t type)
{
	switch (type->u.primitive) {
		case AST_INT:
			return lookup_type(symbol("int"), type->pos);
		case AST_FLOAT:
			return lookup_type(symbol("float"), type->pos);
		case AST_CHAR:
			return lookup_type(symbol("char"), type->pos);
		case AST_STRING:
			return lookup_type(symbol("string"), type->pos);
		case AST_BOOLEAN:
			return lookup_type(symbol("boolean"), type->pos);
		default:
			em_error(type->pos, "Wrong Primitive Type");
			return NULL;
	}
}

// add a new class type into _tenv
// TODO
// class NOT allowed to contain self class type ??
static type_t trans_class_type(ast_type_t type)
{
  return lookup_type(type->u.class_, type->pos);
}

static type_t trans_array_type(ast_type_t type)
{
	switch (type->u.array) {
		case AST_INT:
			return lookup_type(symbol("int[]"), type->pos);
		case AST_FLOAT:
			return lookup_type(symbol("float[]"), type->pos);
		case AST_CHAR:
			return lookup_type(symbol("char[]"), type->pos);
		case AST_STRING:
			return lookup_type(symbol("string[]"), type->pos);
		case AST_BOOLEAN:
			return lookup_type(symbol("boolean[]"), type->pos);
		default:
			em_error(type->pos, "Wrong Array Type");
			return NULL;
	}
}

typedef type_t (*trans_type_func)(ast_type_t);

static trans_type_func _trans_type_funcs[]
=
{
	trans_primitive_type,
	trans_class_type,
	trans_array_type,
};

static type_t trans_type(ast_type_t type)
{
	return _trans_type_funcs[type->kind](type);
}

void sem_trans_prog(ast_decl_t prog)
{
    _venv = env_base_venv();
    _tenv = env_base_tenv();
    
    // add System type to env
    type_t sys_class = ty_class();
    sym_enter(sys_class->u.methods, symbol("printf"), env_func_entry(list(ty_int(), NULL), ty_void()));
    sym_enter(_tenv, symbol("System"), sys_class);
    sym_enter(_venv, symbol("System"), env_var_entry(sys_class));
    
    // start semantic checking...
    trans_decl(prog);
}
